/*
 * Morse code translator
 */
/*
 * TODO:
 *   - binsearch in enc_line()
 *   - support input from cmdline:
 *       % ./morse_code -d "- . ... - .----" -e "test2"
 */
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <errno.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <getopt.h>
#include <stdbool.h>

#define BUFSIZE 8192
#define MORSE_CODE_MAXLEN 7

typedef struct morse_char {
	char			chr;
	char			str[9];
	struct morse_char	*left;
	struct morse_char	*right;
} morse_char;

static morse_char morse_char_map[] = {
	{ 'a', ".-",		NULL, NULL },
	{ 'b', "-...",		NULL, NULL },
	{ 'c', "-.-.",		NULL, NULL },
	{ 'd', "-..",		NULL, NULL },
	{ 'e', ".",		NULL, NULL },
	{ 'f', "..-.",		NULL, NULL },
	{ 'g', "--.",		NULL, NULL },
	{ 'h', "....",		NULL, NULL },
	{ 'i', "..",		NULL, NULL },
	{ 'j', ".---",		NULL, NULL },
	{ 'k', "-.-",		NULL, NULL },
	{ 'l', ".-..",		NULL, NULL },
	{ 'm', "--",		NULL, NULL },
	{ 'n', "-.",		NULL, NULL },
	{ 'o', "---",		NULL, NULL },
	{ 'p', ".--.",		NULL, NULL },
	{ 'q', "--.-",		NULL, NULL },
	{ 'r', ".-.",		NULL, NULL },
	{ 's', "...",		NULL, NULL },
	{ 't', "-",		NULL, NULL },
	{ 'u', "..-",		NULL, NULL },
	{ 'v', "...-",		NULL, NULL },
	{ 'w', ".--",		NULL, NULL },
	{ 'x', "-..-",		NULL, NULL },
	{ 'y', "-.--",		NULL, NULL },
	{ 'z', "--..",		NULL, NULL },
	{ '0', "-----",		NULL, NULL },
	{ '1', ".----",		NULL, NULL },
	{ '2', "..---",		NULL, NULL },
	{ '3', "...--",		NULL, NULL },
	{ '4', "....-",		NULL, NULL },
	{ '5', ".....",		NULL, NULL },
	{ '6', "-....",		NULL, NULL },
	{ '7', "--...",		NULL, NULL },
	{ '8', "---..",		NULL, NULL },
	{ '9', "----.",		NULL, NULL },
	/*
	 * Special chars (sorted as in ASCII).
	 */
	{ '!', "-.-.--",	NULL, NULL },
	{ '"', ".-..-.",	NULL, NULL },
	{ '$', "...-..-",	NULL, NULL },
	{ '&', ".-...",		NULL, NULL },
	{ '\'',".----.",	NULL, NULL },
	{ '(', "-.--.",		NULL, NULL },
	{ ')', "-.--.-",	NULL, NULL },
	{ '+', ".-.-.",		NULL, NULL },
	{ ',', "--..--",	NULL, NULL },
	{ '-', "-....-",	NULL, NULL },
	{ '.', ".-.-.-",	NULL, NULL },
	{ '/', "-..-.",		NULL, NULL },
	{ ':', "---...",	NULL, NULL },
	{ ';', "-.-.-.",	NULL, NULL },
	{ '=', "-...-",		NULL, NULL },
	{ '?', "..--..",	NULL, NULL },
	{ '@', ".--.-.",	NULL, NULL },
	{ '_', "..--.-",	NULL, NULL },
	/*
	 * The following nodes are needed to build the connected tree
	 */
	{  0 , "----",		NULL, NULL },
	{  0 , "--..-",		NULL, NULL },
	{  0 , "-.-.-",		NULL, NULL },
	{  0 , ".--.-",		NULL, NULL },
	{  0 , ".-.-",		NULL, NULL },
	{  0 , ".-..-",		NULL, NULL },
	{  0 , "..--",		NULL, NULL },
	{  0 , "...-..",	NULL, NULL },
	/*
	 * 2nd run
	 */
	{  0 , "---.",		NULL, NULL },
	{  0 , "..--.",		NULL, NULL },
	{  0 , "...-.",		NULL, NULL }
};

#define MORSE_CHAR_N ((int)(sizeof(morse_char_map) / sizeof(morse_char_map[0])))

/*
 * Normal version
 */
int
binsearch(int what, int* where, int size)
{
	int L = 0,
	    R = size - 1,
	    M;
	while (L < R) {
		M = (L+R)/2;
		if (what > where[M])
			L = M+1;
		else
			R = M;
	}
	return L;
}

void
build_morse_tree(morse_char * root, morse_char * set, int n)
{
	int i;

	for (i = 0 ; i < MORSE_CHAR_N ; ++i) {
		if (!strncmp(root->str, set[i].str, n) && (int) strlen(set[i].str) == n+1) {
			switch (set[i].str[n]) {
			case '.':
				root->left = &set[i];
				if (n < MORSE_CODE_MAXLEN)
					build_morse_tree(root->left, set, n+1);
				break;
			case '-':
			case '_':
				root->right = &set[i];
				if (n < MORSE_CODE_MAXLEN)
					build_morse_tree(root->right, set, n+1);
				break;
			}
		}
	}
	return;
}

/* K&R */
void
treeprint(morse_char * p, FILE* fp)
{
	if (p != NULL) {
		treeprint(p->left, fp);
		fprintf(fp, "%c : %s\n", p->chr ? p->chr : ' ', p->str);
		treeprint(p->right, fp);
	}
}

/*
 * Dear diary,
 * this function takes pointer to a morse encoded string 'line'
 * and prints the decoded result to 'fp'.
 * XXX: This function changes input buffer.
 */
void
dec_line(char * line, FILE * fp, morse_char *mt_root)
{
	register unsigned char j = 0;
	char *ch = NULL;
	morse_char *mt_node = mt_root;

	ch = strtok(line, " \n");
	while (ch && *ch) {
		if (*ch == '/') {
			putc(' ', fp);
			goto next;
		}
		mt_node = mt_root;
		j = 0;
		while (ch[j]) {
			if (ch[j] == '.')
				mt_node = mt_node->left;
			else
				mt_node = mt_node->right;
			++j;
		}
		putc(mt_node->chr, fp);
next:
		ch = strtok(NULL, " \n");
	}
	putc('\n', fp);
}

void
enc_line(register const char *line, FILE *fp, morse_char *mt_root)
{
	register unsigned char i = 36;
	while (*line) {
		if (*line == '\n') {
			putc(*(line++), fp);
			continue;
		} else if (isspace(*line)) {
			fputs("/ ", fp);
		}

		if (isalpha(*line)) {
			fprintf(fp, "%s ",
				 morse_char_map[tolower(*line) - 'a'].str);
		} else if (isdigit(*line)) {
			fprintf(fp, "%s ",
				 morse_char_map[*line - '0' + 26].str);
		} else {
			/*
			 * Now that we've filtered out 36 out of 54
			 * characters we can afford linear search.
			 *
			 * But it's still gay.
			 */
			for (i = 36 ; morse_char_map[i].chr ; ++i) {
				if (*line == morse_char_map[i].chr) {
					fprintf(fp, "%s ",
							morse_char_map[i].str);
					break;
				}
			}
		}
		++line;
	}
}

void
usage(bool requested, char **argv)
{
	(void) fprintf(requested ? stdout : stderr, "%s: USAGE: %s [-e|-d] [string]\n",
			argv[0], argv[0]);
	exit(requested ? EXIT_SUCCESS : EXIT_FAILURE);
}

int
main(int argc, char **argv)
{
	char	*in_buf;
	bool	decode = true;
	FILE*	in_fp = stdin;
	FILE*	out_fp = stdout;
	morse_char mtree_root = { '\t', "", NULL, NULL };
	int opt;
	bool close_in_fp = false,
	     close_out_fp = false,
	     read_input = true;

	build_morse_tree(&mtree_root, morse_char_map, 0);
	/* treeprint(&mtree_root, stdout); */

	while (opt = getopt(argc, argv, "e::d::o:i:"), opt != -1) { /* GNU ext. :( */
		switch (opt) {
		case 'd':
		case 'e':
			if (optarg) {
				in_buf = malloc((size_t) (strlen(optarg) + 2u));
				if (!in_buf) {
					(void)fprintf(stderr, "%s: Out of memory\n",
							argv[0]);
					break;
				}
				strcpy(in_buf, optarg);
				if (opt == 'e')
					enc_line(in_buf, out_fp, &mtree_root);
				else
					dec_line(in_buf, out_fp, &mtree_root);
				free(in_buf);
			} else if (opt == 'e') {
				decode = false;
			}
			read_input = false;
			break;
		case 'o':
			if (!close_out_fp) {
				out_fp = fopen(optarg, "w");
				close_out_fp = true;
			}
			break;
		case 'i':
			if (!close_in_fp) {
				in_fp = fopen(optarg, "r");
				if (! in_fp ) {
					perror(argv[0]);
					(void) fprintf(stderr,
						       "%s: defaulting to stderr\n",
						       argv[0]);
				} else {
					close_in_fp = true;
				}
			}
			break;
		default:
			usage(false, argv);
			break;
		}
	}

	if (read_input) {
		in_buf = malloc(BUFSIZE+1);
		if (in_buf == NULL) {
			(void) fprintf(stderr,
				       "%s: Failed to allocate memory for input buffer (%d bytes)\n",
				       argv[0], BUFSIZE+1);
			exit(EXIT_FAILURE);
		}

		while (fgets(in_buf, BUFSIZE, in_fp)) {
			if (decode)
				dec_line(in_buf, out_fp, &mtree_root);
			else
				enc_line(in_buf, out_fp, &mtree_root);
		}
		free(in_buf);
	}
	if (close_in_fp)
		fclose(in_fp);
	if (close_out_fp)
		fclose(out_fp);

	exit(EXIT_SUCCESS);
	return (0);
}
